package com.strp.mptcp.scheduler;

import com.strp.mptcp.common.DataChunk;
import com.strp.mptcp.common.MPTCPMetrics;
import com.strp.mptcp.congestion.CongestionController;
import com.strp.mptcp.path.PathManager;
import com.strp.mptcp.path.Subflow;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;

import java.io.File;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Stream;

/**
 * Scheduler: 从文件拆分数据块，通过多条子流并结合拥塞控制发送，支持基本重试。
 * <p>
 * 主要功能：
 * 1. 维护 inflight 映射，记录发送后待确认的 chunk
 * 2. 发送后通知 CongestionController
 * 3. 采用环形+拥塞窗口筛选策略选择子流
 * 4. 在 ACK 回调时移除 inflight
 * 5. 支持动态数据块大小调整
 * 6. 实现可靠的重传机制
 * 7. 流量控制
 * 8. 断点续传支持
 * 9. 性能指标收集
 *
 * @author 曹行阳
 * @date 2025/6/12
 */
public class Scheduler implements Runnable {
    private final PathManager pathManager;
    private final CongestionController congestionController;
    private final MPTCPMetrics metrics;
    private final BlockingQueue<FileRegion> taskQueue = new LinkedBlockingQueue<>();
    private final ExecutorService schedulerThread = Executors.newSingleThreadExecutor();
    private final AtomicLong globalSequence = new AtomicLong(0);
    private volatile boolean running = true;

    // 配置参数
    private static final int MIN_CHUNK_SIZE = 4 * 1024;  // 4KB
    private static final int MAX_CHUNK_SIZE = 64 * 1024; // 64KB
    private static final int INITIAL_CHUNK_SIZE = 8 * 1024; // 8KB
    private static final int MAX_RETRIES = 3;
    private static final long RETRY_DELAY_MS = 1000;
    private static final int MAX_INFLIGHT_CHUNKS = 1000;
    private static final long RTT_WINDOW_SIZE = 1000; // 1秒

    // 记录待确认的数据块：seq -> ChunkInfo
    private final Map<Long, ChunkInfo> inflight = new ConcurrentHashMap<>();
    private final Map<Integer, Queue<Long>> pathSequences = new ConcurrentHashMap<>();
    private final Map<Integer, CircularBuffer<Long>> pathRtts = new ConcurrentHashMap<>();
    private long rrCounter = 0;
    private int currentChunkSize = INITIAL_CHUNK_SIZE;

    public Scheduler(PathManager pathManager, CongestionController cc, MPTCPMetrics metrics) {
        this.pathManager = pathManager;
        this.congestionController = cc;
        this.metrics = metrics;
    }

    public void start() {
        schedulerThread.submit(this);
    }

    public void downloadFile(File file, long offset, String transferId) {
        taskQueue.offer(new FileRegion(file, offset, transferId));
    }

    /**
     * 上传文件
     * @param sourceFile 源文件
     * @param targetFile 目标文件
     * @param offset 起始偏移量
     * @param transferId 传输ID
     */
    public void uploadFile(File sourceFile, File targetFile, long offset, String transferId) {
        taskQueue.offer(new FileRegion(sourceFile, offset, transferId));
    }

    /**
     * 上层接收到 ACK 后调用：
     * 移除 inflight，更新 RTT，并通知拥塞控制释放窗口。
     */
    public void onAckReceived(long sequence, int pathId) {
        ChunkInfo chunkInfo = inflight.remove(sequence);
        if (chunkInfo != null) {
            long rtt = System.currentTimeMillis() - chunkInfo.sendTime;
            updatePathRtt(pathId, rtt);
            congestionController.onAck(pathId, chunkInfo.size);
            adjustChunkSize(pathId, rtt);

            // 记录性能指标
            metrics.recordChunkTransferred();
            metrics.recordLatency(pathId, rtt);
            metrics.updateFileProgress(chunkInfo.transferId,
                    chunkInfo.offset + chunkInfo.size,
                    chunkInfo.totalSize);
        }
    }

    private void updatePathRtt(int pathId, long rtt) {
        CircularBuffer<Long> rtts = pathRtts.computeIfAbsent(pathId, k -> new CircularBuffer<>(10));
        rtts.add(rtt);
    }

    private void adjustChunkSize(int pathId, long rtt) {
        CircularBuffer<Long> rtts = pathRtts.get(pathId);
        if (rtts != null && rtts.size() >= 10) {
            double avgRtt = rtts.stream().mapToLong(Long::longValue).average().orElse(0);
            if (avgRtt < 50) { // 如果平均 RTT 小于 50ms，增加块大小
                currentChunkSize = Math.min(currentChunkSize * 2, MAX_CHUNK_SIZE);
            } else if (avgRtt > 200) { // 如果平均 RTT 大于 200ms，减小块大小
                currentChunkSize = Math.max(currentChunkSize / 2, MIN_CHUNK_SIZE);
            }
        }
    }

    @Override
    public void run() {
        while (running) {
            try {
                FileRegion region = taskQueue.poll(100, TimeUnit.MILLISECONDS);
                if (region != null) {
                    processFile(region);
                }
                checkTimeouts();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private void checkTimeouts() {
        long currentTime = System.currentTimeMillis();
        List<Long> timeoutSequences = new ArrayList<>();

        for (Map.Entry<Long, ChunkInfo> entry : inflight.entrySet()) {
            ChunkInfo info = entry.getValue();
            if (currentTime - info.sendTime > info.timeout) {
                timeoutSequences.add(entry.getKey());
            }
        }

        for (Long seq : timeoutSequences) {
            ChunkInfo info = inflight.get(seq);
            if (info != null && info.retries < MAX_RETRIES) {
                retryChunk(seq, info);
            } else {
                inflight.remove(seq);
                // TODO: 通知上层处理失败
            }
        }
    }

    private void retryChunk(long sequence, ChunkInfo info) {
        Subflow path = pathManager.getSubflow(info.pathId);
        if (path != null && path.isActive()) {
            info.retries++;
            info.sendTime = System.currentTimeMillis();
            info.timeout = calculateTimeout(info.pathId);

            // 将 DataChunk 转换为 ByteBuf
            ByteBuf payload = Unpooled.buffer();
            payload.writeInt(info.chunk.header().type().ordinal());
            payload.writeInt(info.chunk.header().pathId());
            payload.writeLong(info.chunk.header().sequence());
            if (info.chunk.payload() != null) {
                payload.writeBytes(info.chunk.payload());
            }

            path.send(payload);

            // 记录重传指标
            metrics.recordRetransmission(info.pathId);
        }
    }

    private long calculateTimeout(int pathId) {
        CircularBuffer<Long> rtts = pathRtts.get(pathId);
        if (rtts != null && !rtts.isEmpty()) {
            double avgRtt = rtts.stream().mapToLong(Long::longValue).average().orElse(0);
            return (long) (avgRtt * 2); // 使用平均 RTT 的 2 倍作为超时时间
        }
        return 5000; // 默认 5 秒超时
    }

    private void processFile(FileRegion region) throws Exception {
        try (RandomAccessFile raf = new RandomAccessFile(region.file, "r");
             FileChannel fileChannel = raf.getChannel()) {

            fileChannel.position(region.offset);
            ByteBuffer buffer = ByteBuffer.allocate(currentChunkSize);
            long totalSize = region.file.length();
            long currentOffset = region.offset;

            while (fileChannel.read(buffer) > 0) {
                buffer.flip();

                // 等待直到有可用的子流和足够的窗口空间
                Subflow bestPath = waitForAvailablePath();
                if (bestPath == null) {
                    throw new RuntimeException("No available paths");
                }

                long seq = globalSequence.getAndIncrement();
                ByteBuf payload = Unpooled.copiedBuffer(buffer);

                // 创建数据包
                ByteBuf data = Unpooled.buffer();
                data.writeInt(DataChunk.ChunkType.DATA.ordinal());
                data.writeInt(bestPath.getPathId());
                data.writeLong(seq);
                data.writeInt(payload.readableBytes());
                data.writeBytes(payload);
                payload.release();

                // 发送并记录 inflight
                bestPath.send(data);
                inflight.put(seq, new ChunkInfo(
                        new DataChunk(
                                new DataChunk.Header(
                                        DataChunk.ChunkType.DATA,
                                        bestPath.getPathId(),
                                        seq,
                                        payload.readableBytes()
                                ),
                                payload
                        ),
                        bestPath.getPathId(),
                        payload.readableBytes(),
                        currentOffset,
                        totalSize,
                        region.transferId
                ));

                // 记录性能指标
                metrics.recordBytesTransferred(bestPath.getPathId(), payload.readableBytes());

                // 通知拥塞控制已发送
                congestionController.onSend(bestPath.getPathId(), payload.readableBytes());

                currentOffset += payload.readableBytes();
                buffer.clear();
            }
        }
    }

    private Subflow waitForAvailablePath() throws InterruptedException {
        Subflow path;
        while ((path = nextAvailablePath()) == null) {
            Thread.sleep(10);
        }
        return path;
    }

    /**
     * 环形 + 拥塞窗口筛选子流
     */
    private Subflow nextAvailablePath() {
        List<Subflow> subs = pathManager.getHealthySubflows();
        int size = subs.size();
        if (size == 0) return null;

        for (int i = 0; i < size; i++) {
            int idx = (int) ((rrCounter + i) % size);
            Subflow sf = subs.get(idx);
            if (congestionController.canSend(sf.getPathId()) &&
                    getInflightCount(sf.getPathId()) < MAX_INFLIGHT_CHUNKS) {
                rrCounter = idx + 1;
                return sf;
            }
        }
        return null;
    }

    private int getInflightCount(int pathId) {
        return (int) inflight.values().stream()
                .filter(info -> info.pathId == pathId)
                .count();
    }

    public void stop() {
        running = false;
        schedulerThread.shutdownNow();
    }

    private record FileRegion(File file, long offset, String transferId) {}

    private static class ChunkInfo {
        final DataChunk chunk;
        final int pathId;
        final int size;
        final long offset;
        final long totalSize;
        final String transferId;
        int retries;
        long sendTime;
        long timeout;

        ChunkInfo(DataChunk chunk, int pathId, int size, long offset,
                  long totalSize, String transferId) {
            this.chunk = chunk;
            this.pathId = pathId;
            this.size = size;
            this.offset = offset;
            this.totalSize = totalSize;
            this.transferId = transferId;
            this.retries = 0;
            this.sendTime = System.currentTimeMillis();
            this.timeout = 5000; // 初始超时时间 5 秒
        }
    }

    private static class CircularBuffer<T> {
        private final T[] buffer;
        private int head;
        private int size;
        private final int capacity;

        @SuppressWarnings("unchecked")
        CircularBuffer(int capacity) {
            this.capacity = capacity;
            this.buffer = (T[]) new Object[capacity];
            this.head = 0;
            this.size = 0;
        }

        public void add(T element) {
            buffer[head] = element;
            head = (head + 1) % capacity;
            if (size < capacity) {
                size++;
            }
        }

        public int size() {
            return size;
        }

        public boolean isEmpty() {
            return size == 0;
        }

        public List<T> toList() {
            List<T> result = new ArrayList<>();
            for (int i = 0; i < size; i++) {
                int idx = (head - size + i + capacity) % capacity;
                result.add(buffer[idx]);
            }
            return result;
        }

        public Stream<T> stream() {
            return toList().stream();
        }
    }
}